#!/usr/bin/env python
import random, sys, getopt

class Tree:         
    data = ""
    children = []

# Tree constructors
def leaf(data):
    tree = Tree()
    tree.data = data
    return tree

def nonleaf(data, children):
    tree = Tree()
    tree.data = data
    tree.children = children
    return tree

class Rule:
    def __init__(self, num, lhs, rhs):
        self.num = num
        self.lhs = lhs
        self.rhs = rhs

    num = 0
    lhs = ""
    rhs = []

def parse_grammar_file(grammar_file):
    grammar = {}
    f = open(grammar_file)
    for line in f:
        line = trim_comment(line)
        tokens = line.split()
        if len(tokens) == 0:
            continue
        if len(tokens) < 2:
            raise "format error: should have at least NUM LHS in a line"
        num = float(tokens[0])
        lhs = tokens[1]
        rhs = tokens[2:]
        rule = Rule(num, lhs, rhs)
        multimap_add(grammar, lhs, rule)
    f.close()
    return grammar

def trim_comment(line):
    pos = line.find("#")
    if pos <> -1:
        line = line[0:pos]
    return line

def multimap_add(d, key, value):
    d[key] = d.get(key,[]) + [value]

def generate(grammar):
    start_sym = "START"
    return generate_tree(grammar, start_sym)

def generate_tree(grammar, sym):
    rules = grammar.get(sym, [])
    if len(rules) == 0:
        return leaf(sym)
    rule = random_pick(rules)
    children = [generate_tree(grammar,s) for s in rule.rhs]
    return nonleaf(sym, children)

# randomly pick a rule according to associated numbers
def random_pick(rules):
    nums = [r.num for r in rules]
    return rules[random_pick_nums(nums)]

def random_pick_nums(nums):
    rand = random.random()
    total = sum(nums)
    partial = 0
    for i, num in enumerate(nums):
        partial += num
        if rand < partial/total:
            return i
    
def print_tree(tree):
    print "(", tree.data,
    for child in tree.children:
        print_tree(child)
    print ")",

def print_string(tree):
    if len(tree.children) == 0:
        print tree.data,
    else:
        for child in tree.children:
            print_string(child)

def print_usage():
    print 'USAGE:', sys.argv[0], '[-t] GRAMMAR_FILE [NUM_SENT]'
    
is_print_tree = False
num_sent = 1
try:
    opts, args = getopt.getopt(sys.argv[1:],"ht")
except getopt.GetoptError:
    print_usage()
    sys.exit(2)
for opt, arg in opts:
    if opt == '-h':
        print_usage()
        sys.exit()
    elif opt in ("-t"):
        is_print_tree = True
if len(args) < 1:
    print_usage()
    sys.exit(2)
grammar_file = args[0]
if len(args) >= 2:
    num_sent = int(args[1])

#grammar_file = "grammar"
#num_sent = 2
#is_print_tree = True

grammar = parse_grammar_file(grammar_file)
trees = [generate(grammar) for i in range(num_sent)]
for tree in trees:
    if is_print_tree:
        print_tree(tree)
    else:
        print_string(tree)
    print
    print


    

